Add LTX-2.5 DFR pipeline (keyframe slots, spatial detailing, tiled temporal rounds) - #14567
Add LTX-2.5 DFR pipeline (keyframe slots, spatial detailing, tiled temporal rounds)#14567alexanderar wants to merge 4 commits into
Conversation
Ports DFRPipeline from the Lightricks reference. Stage 1 generates video plus single-pixel-frame keyframe slots at a fraction of the requested resolution on a VAE-aligned segment grid; both are spatially latent-upsampled and stage 2 re-denoises at twice that resolution with the slots re-attached and an optional spatial detailing IC-LoRA active for that stage only. Optional temporal x2/x4 refine rounds tile the canvas at shared keyframes and densify with ancestral Euler. With spatial_upscalings=2 a full-resolution detailing epilogue follows the rounds. The transformer already stored keyframes_abs_pos_embedding for load/save; this wires it into the forward through a new video_keyframes_mask argument, which only a keyframes-aware pipeline passes, so other pipelines are unaffected. The epilogue denoises the whole canvas in one loop and tiles the transformer call inside it, so every Euler step steps a canvas whose tiles have already agreed on their overlaps. Spatial tiles blend under a trapezoidal mask, since neither side of a height or width border holds a known answer. Temporal tiles are cut on the keyframe seams the last refine round stitched on: both windows reproduce a shared keyframe there, so the later one drops its run-up under a rectangular mask rather than averaging it. Conditionings are attached once on the whole canvas and filtered per tile at the token level, and a keyframe two windows share is one token they both read. The epilogue is handed its keyframes rather than asked to generate them. Each carry plane is decoded on its own -- the VAE is causal, so a stacked decode would bleed neighbours -- then Lanczos-stretched x2 in RGB and encoded again at the output resolution, and pinned fully clean. Only the video latent is spatially upsampled. Conditioning fps is snapped to 60 above 30 rather than merely capped there, at every stage. RoPE time is pixel_frame / fps, and the transformer is trained around 24/25/30 and 60; a temporal round taking 24 fps to 48 lands between those, and it shows as stutter at the latent borders. Playback fps is unchanged, so 24 fps with one round still ships 48 fps. A condition's index is read on the canvas num_frames asks for, and the moment it names is carried onto each refine round's longer canvas by scaling its pixel position by 2**round. The scaled position does not generally land on a latent boundary, so it travels as a pixel index rather than through the public latent index; a keyframe conditioning is appended as extra tokens instead of being spliced into the base grid, so it does not need to. height and width must be divisible by 2**spatial_upscalings times the VAE's spatial compression ratio, which makes 4K 3840x2176 rather than 3840x2160. That rule is checked ahead of the looser one every LTX-2 pipeline applies, so the error names the divisor a DFR caller actually has to satisfy. Four details are easy to get wrong, and each is covered by a test after showing up as a visible seam at a tile handover: - The ancestral step injects noise into every token, so the conditioning blend has to be re-applied afterwards. Skipping it lets the strength-0.95 anchor keyframes erode over the schedule, and those anchors are the only thing pinning adjacent tiles onto the same content. - Velocity is converted to x0 with each token's own noise level, not the scalar schedule sigma: a token held at strength s sits at (1 - s) * sigma. - Each tile draws its ancestral noise from a generator seeded seed + 1000 * round + tile, kept separate from the main generator so the draws do not consume state the next tile's initial noising reads. - Two tiles invent the slot that falls in the later one's dropped lead-in. The stitch keeps the earlier tile's frames there, so the earlier tile's copy is the one the canvas holds, and the one the next round must anchor on.
github-actions
Bot
added
documentation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
| temporal_latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained( | ||
| "path/to/converted/temporal_latent_upsampler", torch_dtype=torch.bfloat16 | ||
| ) | ||
| pipe = LTX2DFRPipeline.from_pretrained( |
There was a problem hiding this comment.
I wonder if we can support this with the same pattern as other multi-stage LTX pipelines?
So something like
pipe = LTX2DFRPipeline.from_pretrained(...)
upsample_pipe = LTX2LatentUpsamplePipeline(...)
# stage1
video_latents, audio_latents, keyframes_latents = pipe(..., output_type="latent")
upscaled_video = upsample_pipe(latent=video_latents ....)
upscaled_keyframes = upsample_pipe(latents=keyframes_latents,....)
# stage2
pipe.load_lora_weights(...)
pipe.set_adapters(...)
video_latents, _, keyframes_latents = pipe(latents=upscaled_video, audio_latents=audio_latents, keyframes_latents=upscaled_keyfames, reference_latents=video_latents, ....)
...
# a new temperal_upsample pipeline?
| if ancestral_eta > 0: | ||
| noise = randn_tensor(latents.shape, generator=ancestral_generator, device=device, dtype=torch.float32) | ||
| stepped = ancestral_euler_step( | ||
| latents, denoised, sigma, self.scheduler.sigmas[index + 1], ancestral_eta, noise | ||
| ) |
There was a problem hiding this comment.
I think we should use an existing scheduler here, or implement a new ancestral Euler scheduler if no existing scheduler can be reused.
I also think that implementing a separate temporal upsampling pipeline as suggested in #14567 (comment) would be cleaner here as we could naturally set its scheduler to be an ancestral Euler scheduler instead of needing to switch the scheduling logic mid-pipeline.
There was a problem hiding this comment.
@dg845 the question is how truthful to the reference pipeline in the LTX-2 repo you want to be.
My goal was to be as close as possible to the source. This is exactly what is done in the LTX. The more deviations from the the source we have the more difficult it will be to maintain the compatibility and make sure that diffusers pipeline produces results of the same quality
There was a problem hiding this comment.
Hi @alexanderar, my thoughts are that the pipeline should be as close as possible semantically to the reference, while still respecting the diffusers design. For example, almost all current diffusers pipelines have either a single or nested denoising loop by design. I don't think the above necessarily rules out having a single monolithic pipeline like LTX2DFRPipeline that covers all stages, each with their own denoising loops, but splitting up the pipeline would fit the diffusers design better (including the way LTX-2.0/2.3 two-stage inference is currently implemented). It would also have the advantage of making each stage independently runnable: if e.g. the temporal refinement logic was in a separate pipeline, and I run Stage 1 + spatial upsampling + Stage 2 (the current LTX2DFRPipeline default behavior), then I can potentially run temporal refinement on the Stage 2 outputs later using its pipeline without needing to rerun the previous stages.
@yiyixuxu WDYT?
There was a problem hiding this comment.
Sounds good.
I am working on the requested changes
Addresses the review on huggingface#14567: use the same compose pattern as the other LTX two-stage pipelines, get ancestral Euler from an existing scheduler, and put the temporal rounds in their own pipeline so the schedule is not switched mid-call. `LTX2DFRPipeline.__call__` is now one denoise pass at `height` x `width`. Callers compose stage 1, `LTX2LatentUpsamplePipeline`, stage 2 and each temporal round, and a documented recipe is the copy-paste 1080p path. The recipe knobs (`spatial_upscalings`, `temporal_upscalings`, `detailing_lora_adapter_name`) and the required upsampler components are gone; `height`/`width` are this pass, not the final output. `ancestral_euler_step` is replaced by `LTXEulerAncestralRFScheduler.step` plus a re-application of the conditioning blend, which ancestral noise would otherwise erode on the strength-0.95 seam anchors. The new `LTX2DFRTemporalRefinePipeline` owns that scheduler and one round; stage 1, stage 2 and the epilogue stay on `FlowMatchEulerDiscreteScheduler`. It refuses any other scheduler rather than silently taking a deterministic step and returning a softer canvas. Pack/unpack, `prepare_latents`, `denoise` and `encode_conditions` move to `LTX2DFRCoreMixin` so neither pipeline subclasses the other. Public latents are raw on both sides of every boundary, `output_type="latent"` returns the untrimmed canvas so a slot on the pad is not dropped, and `trim_canvas` does the trim before decode. `LTX2DFRPipelineOutput` carries `keyframes` and `keyframe_positions`, which cannot be re-derived after a round. Verified against the pre-split implementation: bit-exact on dummy components in fp32 and bf16, with and without the IC-LoRA reference, over one and two rounds; and within one bf16 ulp on the real checkpoint, where the only difference is that the split normalizes upsampled latents in fp32 rather than bf16. Also in this pass: - `__call__` takes `video_tiles` (the `epilogue_tiles` layout) instead of a resolved token plan. Resolving one needs the RoPE coordinates that only exist once `prepare_latents` has run, so a caller could not build the plan at all. - `rebuild_epilogue_keyframes` is public and takes and returns raw latents. The composed epilogue needs it, so it was public API in practice while named private, and its normalized return forced callers into `_denormalize_latents`. - Drop prompt enhancement and `num_videos_per_prompt` from the temporal pipeline. Enhancement belongs to stage 1 -- re-running it would denoise the canvas under a different prompt than the one that generated it -- and the batch is set by the incoming latent canvas, so `num_videos_per_prompt > 1` only ever raised. - Fix the docs recipe: `requested_frames` counted latent frames where `trim_canvas` wants pixel frames, truncating a 241-frame render to 25; the three pipelines share components, so place them together instead of offloading one and leaving `temporal_latent_upsampler` off the device; and the detailing IC-LoRA is applied at 0.5, the strength the reference hardcodes. - Pass `crf=0` on the conditions in the refine-round test. The default CRF sends the image through H.264 re-compression, which needs PyAV, so the test failed on any environment without it while testing nothing about re-compression.
|
@dg845 @yiyixuxu |
|
We could potentially mitigate the concern about the recipe only living in the docs as follows (while keeping the split pipeline design): There are a few pipelines which contain sub-pipelines (an example is I think a modular pipeline could more idiomatically support the entire multi-stage workflow (for example, as a combined multi-stage blockset). Linoy has started work on a modular version of the DFR pipeline at #14600. |
|
@alexanderar will also work will @linoytsaban to add similar support into her PR too for DFR |
yiyixuxu
left a comment
yiyixuxu
left a comment
There was a problem hiding this comment.
Thanks so much for working on this, I left one request, I will merge once that's in
Sorry, this is not really ideal, our standard pipeline API is meant for simple task-based inference and cannot have good support for multi-stage inference like LTX. We will aim to have very good support from Modular. and for future releases, hopefully we can make LTX modular-only. so the goal for this PR is to get it in quickly and more or less consistent with current pattern how we support LTX2 in the past
| return latents[:, :, :keep] | ||
|
|
||
|
|
||
| class LTX2DFRCoreMixin: |
There was a problem hiding this comment.
can you remove the mixin and use the #Copied from for common methods shared between pipelines? so that each pipeline is self-contained
4 participants
What does this PR do?
Adds
LTX2DFRPipeline— Diffusion Fidelity Rendering for LTX-2.5 — ported from the Lightricksreference implementation.
Stage 1 generates video plus extra single-pixel-frame keyframe slots at a fraction of the requested
resolution, on a segment grid aligned to the VAE's temporal border. Slots relax the effective temporal
compression at those positions, so the surrounding video can be conditioned on genuinely new frames rather
than interpolated ones. The half-resolution result is kept as an IC-LoRA reference while both the video and
the slot keyframes are upsampled in latent space; stage 2 re-denoises at twice that resolution with the
slots re-attached and an optional x2 spatial detailing IC-LoRA active for that stage only.
Two optional knobs extend it:
temporal_upscalings(0–2) — each round doubles the frame rate: the canvas is temporally upsampled,split into
2 ** roundtiles that meet at shared keyframes, given fresh mid-segment slots, and densifiedwith ancestral Euler. Each tile cross-attends to the slice of the frozen stage-1 audio covering its own
playback window, so both sides of a seam densify against the same sound. Requires the optional
temporal_latent_upsamplercomponent.spatial_upscalings(1 or 2) —2starts the base canvas one more factor of two down and adds afull-resolution detailing pass after the temporal rounds. That pass does not fit in one sequence, so it
denoises the whole canvas in a single loop and tiles the transformer call inside it, which means every
Euler step steps a canvas whose tiles have already agreed on their overlaps. Spatial tiles blend under a
trapezoidal mask; temporal tiles are cut on the keyframe seams the last refine round stitched on.
Whatever padding the canvas needs internally, the caller always gets
(num_frames - 1) * 2 ** temporal_upscalings + 1frames back.Notable details
keyframes_abs_pos_embeddingwas already stored for load/save but never consumed. Thiswires it into the forward through a new optional
video_keyframes_maskargument, which only akeyframes-aware pipeline passes — other pipelines are unaffected, and the default config leaves the branch
inert.
pixel_frame / fps, and the transformer is trained around 24/25/30 and 60; a temporal round taking 24 fpsto 48 lands between those and shows as stutter at the latent borders. Playback fps is unchanged.
height/widthmust be divisible by2 ** spatial_upscalingstimes the VAE's spatialcompression ratio — 64 at the default, 128 at
spatial_upscalings=2. So a 4K run is 3840x2176, not3840x2160. That rule is checked ahead of the generic one so the error names the divisor a DFR caller
actually has to satisfy.
eta=0.5, which is notexpressible through
FlowMatchEulerDiscreteScheduler— itsstochastic_samplingbranch renoises fullyfrom
x0with noeta, no intermediatesigma_down, and no variance-preserving rescale, so it differseven at
eta=1.ancestral_euler_stepis a module-level function with that reasoning in its docstring.pipeline_ltx2_dfr.py(pipeline) anddfr_layout.py(canvas layout: segment grid, tileplans, blend masks, token plan). The conversion script gains
--temporal_latent_upsamplerfor the x2temporal latent upsampler, which is not part of the base repo.
Tests
Pipeline-level (
tests/pipelines/ltx2/test_pipeline_ltx2_dfr.py) and layout-level(
test_ltx2_dfr_layout.py), plus two additions to the existing transformer model tests for the keyframeembedding. Following
.ai/references/testing.md: pytest-style config class +PipelineTesterMixin/MemoryTesterMixinonly, real components at tiny config, no LoRA or@slowtests in this first pass.mainbasemake quality(includingutils/check_ai.py),utils/check_copies.py,utils/check_dummies.pycleanUpstream note
Porting this surfaced a bug in the reference implementation: an image conditioning with a non-zero frame
index is mis-placed (and usually dropped) after a temporal round, because its position is never scaled onto
the round's longer canvas. Reported to the Lightricks team, confirmed, and being fixed upstream. This
pipeline scales in both the refine rounds and the epilogue, and the placement is pinned by a test.
Before submitting
self-reviewskill on the diff?documentation guidelines, and
here are tips on formatting docstrings.
Self-review notes
Ran the
self-reviewskill over four rounds; final verdict READY, no blocking issues. Nine findings wereraised and resolved across rounds 1–4:
denoisepath (video_onlywas alwaysFalse)>= 0selfdfr_layout/ module level_audio_latents_for_tilereturned a count nothing usedowned_segment_countshad a single callersplit_canvas_at_seamstest_ltx2_dfr_layout.pydfr_layoutLTX2DFREpilogueTile/LTX2DFRTokenPlan, matching the module's existing NamedTuple idiomper_token_sigmaread as scheduler duplicationper_token_timestepspath is not equivalentFindings I deliberately did not fix
choose_segment_lengthhas a single caller, which the coding-style guide suggests inlining. Kept because itmirrors a function of the same name in the reference implementation, which keeps future port diffs
readable. The contrast is
owned_segment_counts, which was inlined precisely because it had noreference counterpart.
video_tile_planlives indfr_layoutbut its tests live in the pipeline test file. They need realRoPE coordinates from
prepare_latents; moving them would drag a pipeline fixture into a pure-layout testfile. Placement follows the dependency.
prepare_latents/denoiseto reach internal state. They assert on whatthose methods returned (real RoPE coords compared against
transformer.rope.prepare_video_coords, realtensors across two real passes), not on the arguments passed in. There is no cheaper behavioural proxy for
cross-pass data flow at dummy resolution.
ancestral_euler_stepis a hand-rolled sampler step. Verified against the scheduler rather thanassumed — see "Sampler" above. Flagging it so a reviewer sees the verification instead of re-deriving it.
Deliberate departures from the reference
LTX2VideoCondition.indexis a latent index, per diffusers convention (the reference uses a pixelindex). Positions are carried in pixel space internally so a scaled position is never floored onto the
latent grid.
0.5, and the reference hardcodes that. Thispipeline documents it instead of enforcing it:
set_adapterswrites a scale for every adapter itnames (
Noneresolves to1.0) and there is no API to read a weight back, so enforcing it would silentlyclobber the caller's other adapters. Happy to change this if maintainers prefer — the alternative is
reaching into peft's
BaseTunerLayer.set_scalefor the single adapter.Not ported
The reference recently made its DFR decode keyframe-aware (passing the encoded keyframe planes to the
DiffVAE decoder). That is blocked here: neither
AutoencoderKLLTX2Videonorltx2_diffusion_decoder.pyhasa keyframe path, so it needs the reference's dual-stream / joint-neighbourhood-attention decoder ported into
the diffusers VAE first. That is models-level work, better as its own PR. Everything else from the
reference's current
mainis included.Fidelity checks
Since this is a port, correctness was checked against the reference rather than only by unit test:
dfr_layoutis bit-exact against the reference's tiling and layout modules over ~1200 cases (segmentgrid, seam splits, count splits, both mask shapes, full temporal tile plans).
count and tile offset — including the keep/drop decisions.
0.95, ancestraleta 0.5, fps snap60/30, epiloguespatial overlap
12, epilogue keyframe strength1.0.stage_2_sigmas[0]),temporal tiles and the epilogue frozen at sigma 0.
Two gotchas worth writing down (proposal, not in this diff)
Both came out of this port and apply to pipelines generally, so they are not included here — this PR stays
scoped to the pipeline. Happy to send either as its own small PR against
.ai/references/pipelines.mdifyou think they are worth recording:
set_adapterswrites a scale for every adapter it names, andweights=Noneresolves to1.0, sousing it to pin one adapter's strength silently resets every other active adapter — and there is no API to
read a weight back and preserve it.
set_adapter(singular) activates without touching scaling.position through a latent-index API floors it onto the latent grid, which bites whenever the position is
not a multiple of the temporal scale.
Who can review?
@yiyixuxu @dg845